import simulation
from madin import *

device = 'cpu'
madin = [1, 2, 3, 4, 6, 8, 11, 15, 20, 26, 34, 45, 59, 77, 101, 132, 172, 224, 292, 380, 494, 643, 836, 1087, 1170,
         1254, 1304, 1414, 1521, 1573, 1630, 1696, 1839, 1912, 1979, 2121, 2206, 2391, 2574, 3347, 3825]

madin = [1, 2, 3, 5, 7, 11, 17, 25, 38, 57, 86, 129, 194, 291, 437, 656, 985, 1000, 1000, 1000, 1000, 1000, 1000, 1000,
         1000, 1000, 1000, 1000, 1000, 1000]
# madin = [0] * 10 + [494, 643, 836, 1087, 1170, 1254, 1304, 1414, 1521, 1573, 1630, 1696, 1839, 1912, 1979, 2121, 2206, 2391, 2574, 3347, 3825]
print(madin)


def getrnn(datalist):
    trandata = []
    for i in range(len(datalist) - autocode_torch.span):
        line = []
        for j in datalist[i + 1:i + 1 + autocode_torch.span]:
            line.append(j["Close"])
        trandata.append(line)
    trandata = np.array(trandata, dtype=np.float32)
    trandata = autocode_torch.processInput(trandata)
    trant = torch.from_numpy(trandata).to(device)

    autoencoder = torch.load(autocode_torch.torchfilename).to(device)
    encoded, decoded = autoencoder(trant)
    disres = decoded.data.cpu().numpy()
    disres = disres[:, -1]
    get_EMA(disres, 0.9)
    return disres


def run(context):
    datalist = context.data[-350:]
    ar = getrnn(datalist)
    # if ar[-1] < -0.8:
    #     print(ar[-1])

    context.data[-1]["rnn"] = ar
    opsize = context.baseblance() * 1e-8 / context.point() * context.leverage() * madin[int(len(context.litera()) / 1)]
    df = 0.8
    df2 = 0.
    madinprice = 80
    global madinnum
    if len(context.litera()) > madinnum:
        madinnum = len(context.litera())

    if "rnn" in context.data[-2]:
        if ar[-1] < -df and context.maxbuy() > opsize:
            if (len(context.litera()) == 0) or (
                    len(context.litera()) > 0 and context.litera()[-1]['time'] + 60 * context.qt_type() <=
                    context.data[-1]['Date'] and
                    context.litera()[-1]['prise'] - madinprice * context.point() > context.data[-1]['Close']):
                context.buy(opsize, context.data[-1]["Close"] - 2, context.data[-1]["Close"] + 600 * context.point())
                # print('buy')
        # if ar[-1] > df and context.maxbuy() > opsize:
        #     if (len(context.litera()) == 0) or (
        #             len(context.litera()) > 0 and context.litera()[-1]['time'] + 60 * context.qt_type() <=
        #             context.data[-1]['Date'] and
        #             context.litera()[-1]['prise'] + madinprice * context.point() < context.data[-1]['Close']):
        #         context.sell(opsize, context.data[-1]["Close"] + 2, context.data[-1]["Close"] - 600 * context.point())

        if ar[-1] < df2:
            context.closesell()
        if ar[-1] > -df2:
            context.closebuy()


def getv(one, key):
    if key in one:
        return one[key]
    else:
        return float("NaN")


def DrawIndex(ax, datalist):
    poslist = [one["pos"] for one in datalist]
    disres1 = np.array([getv(one, "buy") for one in datalist])
    disres2 = np.array([getv(one, "sell") for one in datalist])
    # disres3 =np.array([getv(one,"rnn") for one in datalist])
    ax.plot(poslist, disres1, color="r")
    ax.plot(poslist, disres2, color="g")
    ax.hlines(0, poslist[0], poslist[-1], colors='b')


def kk(symbol, qt_type):
    global madinnum
    print(symbol, )
    # datalist = json.load(open("data/%s.json" % symbol, 'r'))
    datalist = finddb(dbname="stock", tablename=symbol, limit=20000)
    datalist = joinKline(datalist, 3)
    # print(len(datalist))

    ch = simulation.ContextHost(datalist[-100000:], blance=10000000)
    ch.product = symbol + 'A1'
    ch.qt_type = qt_type
    ch.calcgap = 250
    ch.leverage = 100
    ch.point = 1e-3 if "JPY" in symbol else 1e-5
    ch.point = 0.01
    ch.spread = 5 * ch.point
    ch.doLoop(run)
    print('madinnum: ', madinnum)
    simulation.goShow(ch, DrawIndex)
    madinnum = 0


if __name__ == "__main__":
    madinnum = 0
    qts = [5, 15]
    qts = [15]
    # sbs = ["USDJPY", "EURUSD", "EURJPY", "GBPUSD"]
    sbs = ["FXUSDJPY", "FXEURUSD", "FXAUDUSD", "FXEURJPY", "FXAUDJPY", "FXEURGBP", "FXGBPUSD", "FXGBPJPY"]
    sbs = ["sz50_5", "zz500_5", "hs300_5"]
    # sbs = ["FXUSDJPY"]
    for qt_type in qts:
        for symbol in sbs:
            kk('' + symbol, qt_type)
            #     kk('N' + symbol, qt_type)
            print(symbol, qt_type, '********************************************************')
